package proxy

import (
	"fmt"
	"io"
	"net"
	"sync"

	"openSesame/internal/whitelist"
	"openSesame/pkg/config"
	"openSesame/pkg/logger"
)

// Server 代理服务器
type Server struct {
	name       string
	listenPort string
	targetPort string
	whitelist  *whitelist.Whitelist
	listener   net.Listener
	stopChan   chan struct{}
	wg         sync.WaitGroup
}

// NewServer 创建新的代理服务器
func NewServer(name, listenPort, targetPort string, whitelist *whitelist.Whitelist) *Server {
	return &Server{
		name:       name,
		listenPort: listenPort,
		targetPort: targetPort,
		whitelist:  whitelist,
		stopChan:   make(chan struct{}),
	}
}

// Start 启动代理服务器
func (s *Server) Start() error {
	// 创建监听器
	listener, err := net.Listen("tcp", s.listenPort)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %v", s.listenPort, err)
	}

	// 创建白名单监听器
	s.listener = whitelist.NewWhitelistListener(listener, s.whitelist)

	logger.Infof("Proxy server %s started on %s", s.name, s.listenPort)

	// 启动接受连接的协程
	s.wg.Add(1)
	go s.acceptLoop()

	return nil
}

// Stop 停止代理服务器
func (s *Server) Stop() {
	close(s.stopChan)
	if s.listener != nil {
		s.listener.Close()
	}
	s.wg.Wait()
	logger.Infof("Proxy server %s stopped", s.name)
}

// acceptLoop 接受连接循环
func (s *Server) acceptLoop() {
	defer s.wg.Done()

	for {
		select {
		case <-s.stopChan:
			return
		default:
			conn, err := s.listener.Accept()
			if err != nil {
				if !isClosedError(err) {
					logger.Errorf("Failed to accept connection: %v", err)
				}
				continue
			}

			// 获取客户端IP
			clientIP := conn.RemoteAddr().(*net.TCPAddr).IP.String()

			// 通知白名单更新访问时间
			s.whitelist.NotifyAccess(clientIP)

			// 处理连接
			s.wg.Add(1)
			go s.handleConnection(conn)
		}
	}
}

// handleConnection 处理单个连接
func (s *Server) handleConnection(clientConn net.Conn) {
	defer s.wg.Done()
	defer clientConn.Close()

	// 连接目标服务器
	targetConn, err := net.Dial("tcp", s.targetPort)
	if err != nil {
		logger.Errorf("Failed to connect to target %s: %v", s.targetPort, err)
		return
	}
	defer targetConn.Close()

	// 双向转发数据
	var wg sync.WaitGroup
	wg.Add(2)

	// 客户端 -> 目标服务器
	go func() {
		defer wg.Done()
		io.Copy(targetConn, clientConn)
		targetConn.(*net.TCPConn).CloseWrite()
	}()

	// 目标服务器 -> 客户端
	go func() {
		defer wg.Done()
		io.Copy(clientConn, targetConn)
		clientConn.(*net.TCPConn).CloseWrite()
	}()

	wg.Wait()
}

// isClosedError 检查错误是否为连接关闭错误
func isClosedError(err error) bool {
	if err == nil {
		return false
	}
	return err.Error() == "use of closed network connection"
}

// Manager 代理服务器管理器
type Manager struct {
	servers map[string]*Server
	mu      sync.RWMutex
}

// NewManager 创建新的代理服务器管理器
func NewManager() *Manager {
	return &Manager{
		servers: make(map[string]*Server),
	}
}

// AddServer 添加代理服务器
func (m *Manager) AddServer(whitelist *whitelist.Whitelist, config *config.ProxyConfig) {
	m.mu.Lock()
	defer m.mu.Unlock()

	server := NewServer(config.Name, config.ListenPort, config.TargetPort, whitelist)
	m.servers[config.Name] = server
}

// StartAll 启动所有代理服务器
func (m *Manager) StartAll() {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, server := range m.servers {
		go func(s *Server) {
			if err := s.Start(); err != nil {
				logger.Errorf("Failed to start proxy server %s: %v", s.name, err)
			}
		}(server)
	}
}

// StopAll 停止所有代理服务器
func (m *Manager) StopAll() {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, server := range m.servers {
		server.Stop()
	}
}
